Lazy neighborhood filter kernel compilation - #1708
Merged
Merged
Conversation
cmdupuis3
added
bug
ASV BenchmarkingBenchmark Comparison ResultsBenchmarks that have improved:
Benchmarks that have stayed the same:
Benchmarks that have got worse:
|
Sevans711
requested changes
Aug 26, 2026
Sevans711
left a comment
Sevans711
left a comment
Collaborator
There was a problem hiding this comment.
Thank you looking into this! Assuming you wanted me to give an initial review now because you requested review, but not 100% sure since it is still marked as draft.
The regression test is clever, and I confirmed it crashes on main. This should be one good way to guard against import uxarray slowdowns in the future!
I left inline comments about the caching. Basically, it feels strange to me to use functools.cache in this situation, and I think it leads to some more-complicated-than-necessary code logic at the kernel call sites. I think it would be easier to maintain/understand quickly in the future if using a single cache which directly contains the cached kernels themselves, instead of caching on a per-function basis. That might also help with introspection/debugging later if you ever want to quickly check which kernels have been compiled so far. For example, there could even be a test which asserts _NUMBA_KERNELS is empty after import uxarray.
cmdupuis3
marked this pull request as ready for review
August 26, 2026 19:16
Collaborator
Author
|
pre-commit.ci autofix |
Sevans711
approved these changes
Aug 31, 2026
Sevans711
left a comment
Sevans711
left a comment
Collaborator
There was a problem hiding this comment.
After the factory pattern rework in defa8d6, I am happy with how this looks. The _make_kernel function itself looks a bit complicated but the docstring and comments clearly explain what is happening, and now it actually returns "a kernel function" instead of returning "a kernel factory which must be called to get the kernel function" so the way to use it is aligned with my intuitions based on the function names and documentation. I recommend merging to main!
erogluorhan
self-requested a review
August 31, 2026 21:19
erogluorhan
approved these changes
Aug 31, 2026
erogluorhan
left a comment
erogluorhan
left a comment
Member
There was a problem hiding this comment.
This looks good to me; thanks for the fixes!
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Closes #1706
Overview
This PR should solve an issue introduced with PR #941, wherein the neighborhood filter gufunc-based kernels would compile eagerly, forcing a major slowdown of basic
import uxarrayimports, as well as creating threadpools that would eventually block safe forking that is needed in PR #1700.The idea is to have lazy compilation of gufuncs introduced with neighborhood filters for performant reductions. This prevents the otherwise eager compilation of every one of these kernels upon
import uxarray.Combined, this PR and PR #1700 can bring the total benchmark suit runtime down to about half of the pre-neighborhood filters runtime of ~40min, yielding about 20min without neighborhood filter benchmarks and 26-40ish minutes with them.
On its own, this PR can at least mitigate some of the benchmark performance regression, on the order of about 30-50% or so.
PR Checklist
General
Testing & Benchmarking
Documentation
AI Disclosure
AI Usage: Claude Opus 5